package com.heima.search.model;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

public class Trie {
    private List<String> list = new ArrayList<>();
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
        root.var = ' ';
    }

    /**
     * 插入trie树
     *
     * @param word
     */
    public void insert(String word) {
        TrieNode ws = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!ws.children.keySet().contains(c)) {
                ws.children.put(c, new TrieNode(c));
            }
            ws = ws.children.get(c);
        }
        ws.isWord = true;
    }

    public void insert(Collection<String> list) {
        for (String s : list) {
            this.insert(s);
        }
    }

    /**
     * 查询trie树
     *
     * @param prefix
     * @return
     */
    public List<String> startWith(String prefix) {
        List<String> match = new ArrayList<>();
        TrieNode ws = root;
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (!ws.children.keySet().contains(c)) return match;
            ws = ws.children.get(c);
            if (!ws.containLongTail) {
                for (char cc : ws.children.keySet()) {
                    match.add(prefix + cc);
                }
            } else {
                //包含长尾词 从map中取
            }
        }
        return match;
    }

    private TrieNode getMyTrieNode(TrieNode ws, String prefix) {
        if (prefix.length() == 0) {
            return ws;
        }
        char c = prefix.charAt(0);
        if (ws.children.keySet().contains(c)) {
            TrieNode node = ws.children.get(c);
            return getMyTrieNode(node, prefix.substring(1));
        }
        return null;
    }

    private void getMyWords(TrieNode ws, String prefix) {
        if (ws == null) {
            return;
        }
        if (ws.isWord) {
            if (list.size() >=10) {
                return;
            }
            list.add(prefix);
        }
        Map<Character, TrieNode> map = ws.children;
        for (Map.Entry<Character, TrieNode> entry : map.entrySet()) {
            Character key = entry.getKey();
            TrieNode value = entry.getValue();
            if (value.children.size() == 0) {
                if (list.size() >=10) {
                    return;
                }
                list.add(prefix + key);
            } else {
                getMyWords(value, prefix + key);
            }
        }
    }


    public static void main(String[] args) {
       Trie t = new Trie();
        t.insert("黑马");
        t.insert("黑白禁区");
        t.insert("黑洞");
        t.insert("黑龙江");
        t.insert("黑龙江的");
        t.insert("黑龙江的哈");
        t.insert("黑龙江的哈尔");
        t.insert("黑龙江的哈尔滨");
        t.insert("黑龙江的哈尔滨市");
        t.insert("黑龙江的哈尔滨市开");
        t.insert("黑龙江的哈尔滨市开通");
        t.insert("黑龙江的哈尔滨市开通了");
        t.insert("黑龙江的哈尔滨市开通了通");
        t.insert("黑龙江的哈尔滨市开通了通上海");
        t.insert("黑龙江的哈尔滨市开通了通北");
        t.insert("黑龙江的哈尔滨市开通了通北京");
        t.insert("黑龙江的哈尔滨市开通了通北京高铁");
        List<String> list1 = t.getWords("黑");
        List<String> list2 = t.getWords("黑龙江的哈尔滨市开通了通");
        // 黑龙江的哈尔滨市开通了通 - children [北]
        // 黑龙江的哈尔滨市开通了通北 - children [京]
        // 黑龙江的哈尔滨市开通了通北京 - children [高]
        // 黑龙江的哈尔滨市开通了通北京高铁 - children []
        List<String> list3 = t.getWords("黑花");
        System.out.println(list1);
        System.out.println(list2);
        System.out.println(list3);
    }

    public List<String> getWords(String prefix) {
        this.list = new ArrayList<>();
        TrieNode node = this.getMyTrieNode(root, prefix);
        this.getMyWords(node, prefix);
        return list;
    }
}